-----------------------------------------------------------------------------
-- Copyright 2012-2021, Microsoft Research, Daan Leijen.
--
-- This is free software; you can redistribute it and/or modify it under the
-- terms of the Apache License, Version 2.0. A copy of the License can be
-- found in the LICENSE file at the root of this distribution.
-----------------------------------------------------------------------------

-- Functions on type variables, including type variable substitution.
-----------------------------------------------------------------------------

module Type.TypeVar
  ( -- * Type substitutable entities
    HasTypeVar (substitute, ftv, btv, ftc),
    (|->),
    alltv,
    fuv,
    fbv,
    fsv,

    -- * Ordered free type variables
    HasOrderedTypeVar,
    oftv,
    ofuv,

    -- * Substitution from type variables to types
    Sub,
    subNull,
    subNew,
    subSingle,
    (@@),
    subCompose,
    subExtend,
    subRemove,
    subIsNull,
    subCount,

    -- ** Breaks abstraction
    subDom,
    subRange,
    subList,
    subCommon,
    subFind,

    -- * Set of type variables
    Tvs,
    tvsNew,
    tvsEmpty,
    tvsIsEmpty,
    tvsSingle,
    tvsInsert,
    tvsInsertAll,
    tvsMember,
    tvsList,
    tvsDiff,
    tvsUnion,
    tvsUnions,
    tvsIntersect,
    tvsCommon,
    tvsRemove,
    tvsFilter,
    tvsDisjoint,
    tvsIsSubsetOf,

    -- * Effects
    posnegEffects,

    -- * Requiring HasUnique
    freshTypeVar,

    -- * Internal
    subInsert,
    subInserts,
    -- Equal types
    matchType,
    -- Type constructors
    TypeCons,
    tcsEmpty,
    tcsUnion,
    tcsUnions,
  )
where

import Lib.Trace

import Common.Failure (HasCallStack, assertion)
import Common.File (seqList)
import Common.Range
import Common.Unique
import Data.List (foldl', nub, partition)
import qualified Data.Map.Strict as M
import qualified Data.Set as S
-- instances

import Kind.Kind
import Kind.Pretty ()
import Type.Kind
import Type.Type
import qualified Data.Maybe as M

{--------------------------------------------------------------------------

--------------------------------------------------------------------------}

-- | The free unifiable (eg 'Meta') type variables
fuv :: (HasTypeVar t) => t -> Tvs
fuv t =
  tvsFilter isMeta (ftv t)

-- | The free 'Skolem' type variables
fsv :: (HasTypeVar t) => t -> Tvs
fsv t =
  tvsFilter isSkolem (ftv t)

-- | The free 'Bound' type variables
fbv :: (HasTypeVar t) => t -> Tvs
fbv t =
  tvsFilter isBound (ftv t)

-- | All type variables (including the bound ones)
alltv :: (HasTypeVar t) => t -> Tvs
alltv t =
  tvsUnion (btv t) (ftv t)

{--------------------------------------------------------------------------
  Substitution

  For a substitution is should hold that:
  (s1@@s2) |-> x  <==>  s1 |-> (s2 -> x)

  We can implement this by:
  1) at composition we apply the first substitution to the other, and
     finding an identifier is a simple lookup.
  or,
  2) we compose by simple union and perform a fixpoint at lookup.

  We have chosen (1) here, but it can be interesting to compare
  performance with strategy (2).
--------------------------------------------------------------------------}

-- | A substitution from type variables to 'Tau's.
data Sub = Sub {subMap :: !(M.Map TypeVar Tau), subTvs :: !Tvs}

unSub (Sub sub _) =
  sub

subCount :: Sub -> Int
subCount (Sub sub _) =
  M.size sub

subNull :: Sub
subNull =
  Sub M.empty tvsEmpty

subIsNull :: Sub -> Bool
subIsNull (Sub sub _) =
  M.null sub

subNew :: (HasCallStack) => [(TypeVar, Tau)] -> Sub
subNew sub =
  -- assertion "Type.TypeVar.subNew" (all (\tv -> length (filter (==tv) tvs) == 1) tvs) $
  -- assertion "Type.TypeVar.subNew.Tau" (all isTau taus) $
  let s =
        assertion
          ( "Type.TypeVar.subNew.KindMismatch: length "
              ++ show (length sub)
              ++ ": "
              ++ unlines (map (\(x, t) -> "(" ++ showTypeVar x ++ " |-> " ++ showTp t ++ ")") sub)
          )
          (all (\(x, t) -> getKind x == getKind t) sub)
          $ Sub
            (M.fromList sub)
            (tvsUnion (tvsNew (map fst sub)) (ftv (map snd sub)))
   in seq s s

subDom :: Sub -> Tvs
subDom (Sub sub _) =
  tvsNew (M.keys sub)

subRange :: Sub -> [Tau]
subRange (Sub sub _) =
  M.elems sub

subList :: Sub -> [(TypeVar, Tau)]
subList (Sub sub _) =
  M.toList sub

subCommon :: Sub -> Sub -> [(TypeVar, (Tau, Tau))]
subCommon (Sub sub1 _) (Sub sub2 _) =
  M.toList (M.intersectionWith (,) sub1 sub2)

subSingle :: (HasCallStack) => TypeVar -> Tau -> Sub
subSingle tvar tau =
  -- Top assertion is invalid; it can happen (and happens) in the CoreF typechecker when
  -- typechecking (forall a. f a) with f :: forall b. b -> b, that a bound variable (b) with
  -- number ID must be substituted for another bound variable (a), which *could* have the same
  -- ID. If we want to avoid this, we must ensure that all IDs are distinct; in particular,
  -- the IDs of built-in types such as .select must be distinct from further IDs generated
  -- by the compiler.
  assertion
    ("Type.TypeVar.subSingle: recursive type: " ++ showTVar tvar)
    (not (tvsMember tvar (ftv tau)))
    $
    -- assertion ("Type.TypeVar.subSingle: not a tau") (isTau tau) $
    assertion "Type.TypeVar.subSingle.KindMismatch" (getKind tvar == getKind tau)
    $ Sub (M.singleton tvar tau) (tvsInsert tvar (ftv tau))

subLookup :: HasCallStack => TypeVar -> Sub -> Maybe Tau
subLookup tvar (Sub sub _) =
  case M.lookupIndex tvar sub of
    Nothing  -> Nothing
    Just i   -> let (keyTvar,tau) = M.elemAt i sub
                in -- TVars with a different flavour should have different id's
                   assertion ("Substitute with different flavour: " ++ show tvar ++ " != " ++ show keyTvar)
                                (typevarFlavour keyTvar == typevarFlavour tvar) $
                   -- kinds should always match
                   assertion ("Type.TypeVar.subFind: incompatible kind: "
                          ++ "\n tvar: " ++ showTVar tvar ++ ":" ++ show (getKind tvar) ++ ","
                          ++ "\n type: " ++ show tau ++ ":" ++ show (getKind tau))
                          (getKind tvar == getKind tau) $
                   Just tau

subRemove :: [TypeVar] -> Sub -> Sub
subRemove tvars (Sub sub tvs) =
  {-
  (if any (\tv -> M.member tv sub)  tvars
    then let subtvs =  filter (\(t,_) -> elem t tvars) (M.toList sub)
         in trace ("forall " ++ show tvars ++ " overlaps " ++ show subtvs)
    else id
  ) $ -}
  Sub (foldr M.delete sub tvars) (tvsRemove tvars tvs)

subFind :: (HasCallStack) => TypeVar -> Sub -> Tau
subFind tvar sub =
  case subLookup tvar sub of
    Nothing  -> TVar tvar
    Just tau -> tau

showTVar (TypeVar id _ _) =
  show id

-- | Compose two substitutions.
-- @(s1 \@\@ s2) |-> x == s1 |-> (s2 |-> x)@
(@@) :: Sub -> Sub -> Sub
sub1 @@ sub2 =
  subCompose sub1 sub2

subCompose :: Sub -> Sub -> Sub
subCompose sub1@(Sub _ tvs1) sub2@(Sub _ tvs2) =
  Sub
    (M.union (unSub sub1) (unSub (sub1 |-> sub2))) -- ASSUME: left biased union
    (tvsUnion tvs1 tvs2)

subExtend :: TypeVar -> Tau -> Sub -> Sub
subExtend tvar tau sub =
  subSingle tvar tau @@ sub

-- | Insert a new substitution. (Note: breaks abstraction barrier).
subInsert :: TypeVar -> Tau -> Sub -> Sub
subInsert tvar tau (Sub s tvs) =
  assertion
    ("Type.TypeVar.subSingle: recursive type: " ++ showTVar tvar)
    (not (tvsMember tvar (ftv tau)))
    $ assertion ("Type.TypeVar.subSingle: not a tau") (isTau tau)
    $ Sub (M.insert tvar tau s) (tvsUnion tvs (tvsInsert tvar (ftv tau)))

subInserts :: [(TypeVar, Tau)] -> Sub -> Sub
subInserts assoc (Sub sub tvs) =
  Sub
    (M.union (M.fromList assoc) sub) -- ASSUME: left-biased union
    (tvsUnions [tvs, tvsNew (map fst assoc), ftv (map snd assoc)])

(|->) :: (HasCallStack, HasTypeVar a) => Sub -> a -> a
sub |-> x =
  if subIsNull sub
    then x
    else
      (sub `substitute` x)

instance HasTypeVar Sub where
  sub@(Sub _ tvs1) `substitute` (Sub s tvs2) =
    Sub (M.map (\k -> sub `substitute` k) s) (tvsUnion tvs1 tvs2)

  ftv (Sub sub _) =
    tvsUnion (tvsNew (M.keys sub)) (ftv (M.elems sub))

  btv sub =
    tvsEmpty

  ftc (Sub sub _) =
    ftc (M.elems sub)

instance (HasTypeVar a) => HasTypeVar (Maybe a) where
  sub `substitute` mb = case mb of
    Just x -> Just (sub `substitute` x)
    Nothing -> Nothing
  ftv mb = case mb of
    Just x -> ftv x
    Nothing -> tvsEmpty
  btv mb = case mb of
    Just x -> btv x
    Nothing -> tvsEmpty
  ftc mb = case mb of
    Just x -> ftc x
    Nothing -> S.empty

instance (HasTypeVar a, HasTypeVar b) => HasTypeVar (Either a b) where
  sub `substitute` lr = case lr of
    Right x -> Right (sub `substitute` x)
    Left y -> Left (sub `substitute` y)
  ftv lr = case lr of
    Right x -> ftv x
    Left y -> ftv y
  btv lr = case lr of
    Right x -> btv x
    Left y -> btv y
  ftc lr = case lr of
    Right x -> ftc x
    Left y -> ftc y

{--------------------------------------------------------------------------
  Type variables
--------------------------------------------------------------------------}

-- | A set of type variables
data Tvs = Tvs !(S.Set TypeVar)

tvsEmpty :: Tvs
tvsEmpty =
  Tvs S.empty

tvsIsEmpty :: Tvs -> Bool
tvsIsEmpty (Tvs set) =
  S.null set

tvsSingle :: TypeVar -> Tvs
tvsSingle tvar =
  Tvs (S.singleton tvar)

tvsInsert :: TypeVar -> Tvs -> Tvs
tvsInsert tvar (Tvs set) =
  {- assertion "Type.TypeVar.tvsInsert" (case S.lookup tvar set of
                                              Nothing    -> True
                                              Just tvar' -> getKind tvar == getKind tvar') $
  -}
  Tvs (S.insert tvar set)

tvsInsertAll :: [TypeVar] -> Tvs -> Tvs
tvsInsertAll vars (Tvs set) =
  Tvs (foldr S.insert set vars)

tvsNew :: [TypeVar] -> Tvs
tvsNew tvars =
  Tvs (S.fromList tvars)

tvsList :: Tvs -> [TypeVar]
tvsList (Tvs set) =
  S.toList set

tvsRemove :: [TypeVar] -> Tvs -> Tvs
tvsRemove tvar (Tvs set) =
  Tvs (foldr S.delete set tvar)

tvsDiff :: Tvs -> Tvs -> Tvs
tvsDiff (Tvs set1) (Tvs set2) =
  Tvs (S.difference set1 set2)

tvsUnion :: Tvs -> Tvs -> Tvs
tvsUnion (Tvs set1) (Tvs set2) =
  Tvs (S.union set1 set2)

tvsFilter :: (TypeVar -> Bool) -> Tvs -> Tvs
tvsFilter p (Tvs set) =
  Tvs (S.filter p set)

tvsUnions :: [Tvs] -> Tvs
tvsUnions xs =
  foldr tvsUnion tvsEmpty xs

tvsMember :: TypeVar -> Tvs -> Bool
tvsMember tvar (Tvs set) =
  S.member tvar set

tvsIntersect :: Tvs -> Tvs -> Tvs
tvsIntersect (Tvs set1) (Tvs set2) =
  Tvs (S.intersection set1 set2)

tvsCommon :: Tvs -> Tvs -> Bool
tvsCommon tvs1 tvs2 =
  not (tvsDisjoint tvs1 tvs2)

tvsDisjoint :: Tvs -> Tvs -> Bool
tvsDisjoint tvs1 tvs2 =
  tvsIsEmpty (tvsIntersect tvs1 tvs2)

tvsIsSubsetOf :: Tvs -> Tvs -> Bool
tvsIsSubsetOf tvs1 tvs2 =
  tvsIsEmpty (tvsFilter (\tvar -> not (tvsMember tvar tvs2)) tvs1)

assocMax :: [(TypeVar, Tau)] -> Int
assocMax assoc =
  foldl' (\n (tv, t) -> max (tvsMax (ftv t)) (max (typevarId tv) n)) 0 assoc

tvarsMax :: [TypeVar] -> Int
tvarsMax tvars =
  foldl' (\n tv -> max (typevarId tv) n) 0 tvars

tvsMax :: Tvs -> Int
tvsMax (Tvs tvs) =
  S.foldl' (\n tv -> max (typevarId tv) n) 0 tvs

{--------------------------------------------------------------------------
  Entities with type variables
--------------------------------------------------------------------------}

-- | Return the free unifiable type variables in a particular order
ofuv :: (HasOrderedTypeVar a) => a -> [TypeVar]
ofuv x = filter isMeta (oftv x)

-- | Return the free type variables in a particular order
oftv :: (HasOrderedTypeVar a) => a -> [TypeVar]
oftv x =
  let vs = nub (odftv x)
      (es, ws) = partition (\v -> isSpecialKind (typevarKind v)) vs
   in ws ++ es
  where
    isSpecialKind kind =
      (kind == kindEffect || kind == kindHeap || kind == kindLabel || kind == kindFun kindLabel kindEffect)

type TypeCons = S.Set TypeCon

tcsEmpty :: TypeCons
tcsEmpty = S.empty

tcsSingleton :: TypeCon -> TypeCons
tcsSingleton = S.singleton

tcsUnion :: TypeCons -> TypeCons -> TypeCons
tcsUnion = S.union

tcsUnions :: [TypeCons] -> TypeCons
tcsUnions = S.unions

-- | Entitities that contain type variables.
class HasTypeVar a where
  -- | Substitute type variables by 'Tau' types
  substitute :: (HasCallStack) => Sub -> a -> a

  -- | Return free type variables
  ftv :: a -> Tvs

  -- | Return bound type variables
  btv :: a -> Tvs

  -- | Return used type constructors
  ftc :: a -> TypeCons

-- | Entities that contain type variables that can be put in a particular order
class HasOrderedTypeVar a where
  -- | Return free type variables in a particular order, may contain duplicates
  odftv :: a -> [TypeVar]

instance (HasTypeVar a, HasTypeVar b, HasTypeVar c, HasTypeVar d, HasTypeVar e) => HasTypeVar (a, b, c, d, e) where
  sub `substitute` (a, b, c, d, e) =
    (sub `substitute` a, sub `substitute` b, sub `substitute` c, sub `substitute` d, sub `substitute` e)
  ftv (a, b, c, d, e) =
    tvsUnions [ftv a, ftv b, ftv c, ftv d, ftv e]
  btv (a, b, c, d, e) =
    tvsUnions [btv a, btv b, btv c, btv d, btv e]
  ftc (a, b, c, d, e) =
    S.unions [ftc a, ftc b, ftc c, ftc d, ftc e]

instance (HasTypeVar a, HasTypeVar b) => HasTypeVar (a, b) where
  sub `substitute` (x, y) =
    (sub `substitute` x, sub `substitute` y)
  ftv (x, y) =
    tvsUnion (ftv x) (ftv y)
  btv (x, y) =
    tvsUnion (btv x) (btv y)
  ftc (x, y) =
    S.union (ftc x) (ftc y)

instance (HasOrderedTypeVar a, HasOrderedTypeVar b) => HasOrderedTypeVar (a, b) where
  odftv (x, y) =
    odftv x ++ odftv y

instance (HasTypeVar a) => HasTypeVar [a] where
  sub `substitute` xs =
    map (sub `substitute`) xs
  ftv xs =
    tvsUnions (map ftv xs)
  btv xs =
    tvsUnions (map btv xs)
  ftc xs =
    S.unions (map ftc xs)

instance HasTypeVar Range where
  sub `substitute` r = r
  ftv r = tvsEmpty
  btv r = tvsEmpty
  ftc r = S.empty

instance (HasOrderedTypeVar a) => HasOrderedTypeVar [a] where
  odftv xs =
    concatMap odftv xs

instance HasTypeVar Type where
  sub `substitute` tp =
    case tp of
      TForall [] tp   -> TForall [] (sub |-> tp)
      TForall vars tp ->
        let sub' = subRemove vars sub
        in -- if True -- all (\tv -> not (tvsMember tv (subTvs sub'))) vars
           --   then
           assertion ("forall captures substitution: " ++ show vars ++ " " ++ show (subList sub'))
                    (all (\tv -> not (tvsMember tv (subTvs sub'))) vars) $
           TForall vars (sub' |-> tp)
           {-
              else
                let uniq = max (tvsMax (subTvs sub')) (tvsMax (ftv tp)) + 2000
                    tvsub = [(tv, tv {typevarId = typevarId tv + uniq}) | tv <- vars]
                    sksub = subNew [(tv, TVar tvnew) | (tv, tvnew) <- tvsub]
                    preds' = (sub' |-> (sksub |-> preds))
                    newvars = map snd tvsub
                 in seqList newvars $
                      seqList preds' $
                        TForall newvars preds' $!
                          (sub' |-> (sksub |-> tp))
           -}
      TFun args effect result ->
        let args' = (map (\(name, tp) -> let tp' = sub `substitute` tp in seq tp' (name, tp')) args)
         in seqList args' $ TFun args' (sub `substitute` effect) (sub `substitute` result)
      TCon tcon -> TCon tcon
      TVar tvar -> subFind tvar sub
      TApp tp arg -> TApp (sub `substitute` tp) (sub `substitute` arg)
      TSyn syn xs tp -> let xs' = sub `substitute` xs in seqList xs' $ TSyn syn xs' (sub `substitute` tp)

  ftv tp =
    case tp of
      TForall vars tp -> tvsRemove vars (ftv tp)
      TFun args effect result -> tvsUnions (ftv effect : ftv result : map (ftv . snd) args)
      TCon tcon -> tvsEmpty
      TVar tvar -> tvsSingle tvar
      TApp tp arg -> tvsUnion (ftv tp) (ftv arg)
      TSyn syn xs tp -> tvsUnion (ftv xs) (ftv tp)

  btv tp =
    case tp of
      TForall vars tp -> tvsInsertAll vars (btv tp)
      TFun args effect result -> tvsUnions (btv effect : btv result : map (btv . snd) args)
      TSyn syn xs tp -> btv tp
      TApp tp arg -> tvsUnion (btv tp) (btv arg)
      _ -> tvsEmpty

  ftc tp =
    case tp of
      TForall vars tp -> ftc tp
      TFun args effect result -> S.unions (ftc effect : ftc result : map (ftc . snd) args)
      TCon tcon -> S.singleton tcon
      TVar tvar -> S.empty
      TApp tp arg -> S.union (ftc tp) (ftc arg)
      TSyn syn xs tp -> S.union (ftc xs) (ftc tp)

instance HasTypeVar Name where
  sub `substitute` name = name
  ftv name = tvsEmpty
  btv name = tvsEmpty
  ftc name = S.empty

instance HasOrderedTypeVar Type where
  odftv tp =
    case tp of
      TForall vars tp -> filter (\tv -> not (tv `elem` vars)) (odftv tp)
      TFun args effect result -> concatMap odftv (map snd args ++ [effect, result])
      TCon tcon -> []
      TVar tvar -> [tvar]
      TApp tp arg -> odftv tp ++ odftv arg
      TSyn syn xs tp -> odftv tp ++ concatMap odftv xs


{--------------------------------------------------------------------------
  Strictly positive and negative type variables
--------------------------------------------------------------------------}
posnegEffects :: Type -> ([TypeVar], [TypeVar])
posnegEffects tp =
  let tvs1 = positive tp
      tvs2 = negative tp
      tvsPos = tvsDiff tvs1 tvs2
      tvsNeg = tvsDiff tvs2 tvs1
   in ( tvsList $ tvsFilter (\tv -> getKind tv == kindEffect) tvsPos,
        tvsList $ tvsFilter (\tv -> getKind tv == kindEffect) tvsNeg
      )

positive, negative :: Type -> Tvs
positive = posneg True
negative = posneg False

posneg :: Bool -> Type -> Tvs
posneg isPos tp =
  case tp of
    TForall vars tp -> tvsRemove vars (posneg isPos tp)
    TFun args effect result -> tvsUnions (posneg isPos effect : posneg isPos result : map (posneg (not isPos) . snd) args)
    TCon tcon -> tvsEmpty
    TVar tvar -> if (isPos) then tvsSingle tvar else tvsEmpty
    TApp tp args -> tvsUnions (posneg isPos tp : map (posneg isPos) args)
    TSyn syn xs tp -> posneg isPos tp

{--------------------------------------------------------------------------
  Operations requiring HasUnique
--------------------------------------------------------------------------}

freshTypeVar :: (HasUnique m) => Kind -> Flavour -> m TypeVar
freshTypeVar kind flavour =
  do
    id <-
      uniqueId
        ( case flavour of
            Meta -> "_v"
            Skolem -> "$v"
            Bound -> "v"
        )
    -- trace ("Type.TypeVar.freshTypeVar: " ++ show id) $
    return $ TypeVar id kind flavour

{--------------------------------------------------------------------------
  Equality between types
--------------------------------------------------------------------------}
matchType :: HasCallStack => Type -> Type -> Bool
matchType tp1 tp2 =
  case (expandSyn tp1, expandSyn tp2) of
    (TForall vs1 t1, TForall vs2 t2) ->
      if (vs1 == vs2)
        then (eqType t1 t2)
        else
          if (length vs1 == length vs2 && all (\(v1, v2) -> getKind v1 == getKind v2) (zip vs1 vs2))
            then
              let sub = subNew (zip vs1 (map TVar vs2))
               in matchType (sub |-> t1) t2
            else False
    (TFun pars1 eff1 t1, TFun pars2 eff2 t2) -> (matchTypes (map snd pars1) (map snd pars2) && matchEffect eff1 eff2 && matchType t1 t2)
    (TCon c1, TCon c2) -> c1 == c2
    (TVar v1, TVar v2) -> v1 == v2
    (TApp t1 ts1, TApp t2 ts2) -> (matchType t1 t2 && matchTypes ts1 ts2)
    -- (TSyn syn1 ts1 t1, TSyn syn2 ts2 t2)      -> (syn1 == syn2 && matchTypes ts1 ts2 && matchType t1 t2)
    _ -> False

matchTypes ts1 ts2 =
  and (zipWith matchType ts1 ts2)

matchEffect eff1 eff2 =
  matchType (orderEffect eff1) (orderEffect eff2)


------------------------------------------------------------------------------
-- This is just for debugging purposes as we cannot include Type.Pretty :-(
------------------------------------------------------------------------------
showTypeVar :: TypeVar -> String
showTypeVar (TypeVar name kind _) =
  show name ++ "::" ++ show kind

showTp :: Type -> String
showTp tp =
  case tp of
    TVar tvar -> showTypeVar tvar
    TCon tcon -> show (typeconName tcon) ++ "::" ++ show (typeconKind tcon)
    TApp tp args -> showTp tp ++ "<" ++ concatMap (\t -> showTp t ++ ",") args ++ ">"
    TSyn syn args body -> "(syn:" ++ show (typesynName syn) ++ "::" ++ show (typesynKind syn) ++ "<" ++ concatMap (\t -> showTp t ++ ",") args ++ ">" ++ "[" ++ showTp body ++ "])"
    _ -> "?"
